import ast
import math
import os
import argparse
from mindspore.train.callback import CheckpointConfig, TimeMonitor, LossMonitor, ModelCheckpoint
from mindspore import context, nn, Tensor, set_seed, Model
from mindspore.context import ParallelMode
from mindspore.communication.management import init, get_rank, get_group_size
from src.dataset import create_dataset_test
from src.lr_generator import get_lr
from src.model import BCNN
import mindspore.common.dtype as mstype
from mindspore.nn import Accuracy, Recall, F1
from mindspore.train.serialization import load_checkpoint, load_param_into_net
from mindspore import ops
from src.dataset import create_dataset

local_input_url = '/home/d1/dataset/cub200/data'
local_output_url = './output'

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_train_images", type=int, default=5994)
    parser.add_argument("--num_test_images", type=int, default=5794)
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--ckpt_path", type=str, default=None)
    parser.add_argument("--device_target", type=str, default="GPU")
    parser.add_argument("--dataset_path", type=str, default='/home/d1/dataset/CUB_200')
    parser.add_argument("--cnn_name", type=str, default="vgg")
    parser.add_argument("--pre_trained", type=str, default=None, choices=[None, "vgg16.ckpt", "resnet101.ckpt", "densenet121.ckpt"])
    config = parser.parse_args()

    num_train_images = config.num_train_images
    batch_size = config.batch_size
    steps_per_epoch = math.ceil(num_train_images / batch_size)
    
    context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
    net = BCNN(200, config.pre_trained, False, config.cnn_name)
    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    
    test_data_set = create_dataset(config.dataset_path,batch_size=batch_size,is_train=False)
    
    # test_data_loader = test_data_set.create_dict_iterator(output_numpy=True)
    # success_num = 0.0
    # total_num = 0.0
    
    # print("test start, ckpt file path is {}".format(config.ckpt_path))
    # load_param_into_net(net, load_checkpoint(config.ckpt_path))
    # for _, data in enumerate(test_data_loader):
    #     image_data = Tensor(data['image'], mstype.float32)
    #     label = Tensor(data["label"], mstype.int32)
    #     result = net(image_data)
    #     # print(result)
    #     # print(result.shape)
    #     result_label, _ = ops.ArgMaxWithValue(1)(result)
    #     success_num = success_num + sum((result_label == label).asnumpy())
    #     total_num = total_num + float(image_data.shape[0])
    #     # print(result_label)
    #     # print(label)
    # print("accuracy is {}".format(str(round(success_num / total_num, 3))))
    
    print("test start, ckpt file path is {}".format(config.ckpt_path))
    load_param_into_net(net, load_checkpoint(config.ckpt_path))
    optimizer = nn.SGD(net.trainable_params(), learning_rate=0.1, momentum=0.9, weight_decay=1e-04)
    model = Model(net, loss_fn=loss, optimizer=optimizer, metrics={"Accuracy": Accuracy()})
    acc = model.eval(test_data_set, dataset_sink_mode=False)
    print(acc)